# 初始化环境
from env.airsim_gym3 import AirSimMultiAgentEnv
import numpy as np
from datetime import datetime
import time
import cv2
import os


def test_airsim_multi_agent_env():
    """测试 AirSimMultiAgentEnv 类的功能并保存运动过程中的图像和雷达数据"""

    try:
        # 创建保存数据的目录
        timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
        save_dir = f"API_Test/airsim_data_{timestamp}"
        os.makedirs(save_dir, exist_ok=True)
        print(f"数据将保存到目录: {save_dir}")

        # 创建包含3个智能体的环境实例
        print("正在创建环境实例...")
        env = AirSimMultiAgentEnv()

        # 目标位置
        target_pos = [20, 5, 0]
        print(f"目标地点设定...{target_pos}")

        # 重置环境
        print("正在重置环境...")
        env.reset()
        print("环境重置成功")
        prev_distances = [info["distance"] for info in env.get_info(target_pos)]  # 记录初始距离
        print(f"初始距离{prev_distances}")

        # 获取并保存初始状态
        print("正在获取初始环境数据...")
        save_environment_data(env, save_dir, step=0)

        # 设置三辆车的大幅度运动动作
        print("\n正在执行大幅度运动并采集数据...")
        actions = [
            [4, 0, 0],  # 智能体1: 高速前进并大幅度右转
            [0.8, 0, 0],  # 智能体2: 高速前进并大幅度左转
            [0.6, 0, 0]  # 智能体3: 高速前进并更大幅度右转
        ]

        # 运动过程中多次采集数据
        total_steps = 50  # 总步数
        step_interval = 0.2  # 每步间隔时间（秒）

        for step in range(1, total_steps + 1):
            # 执行动作
            states, rewards, dones, infos = env.step(actions)
            print(f"\n执行步骤 {step}/{total_steps}")

            print(f"\nStep奖励: {rewards}")
            print(f"\n各车距离目标: {[info['distance'] for info in infos['infos']]}")

            # 打印距离传感器数据
            for i, info in enumerate(infos['infos']):
                distance_sensor_data = info.get('distance_sensor_data', [])
                print(f"智能体 {i + 1} 距离传感器数据: {distance_sensor_data}")

            # 保存当前步骤的环境数据
            save_environment_data(env, save_dir, step)

            # 等待一段时间
            time.sleep(step_interval)

        # 比较初始和最终状态
        compare_initial_final_states(save_dir, total_steps)

        # 测试获取信息
        print("\n正在测试获取信息...")
        target_pos = [20, 10, 0]
        infos = env.get_info(target_pos)
        if infos:
            for i, info in enumerate(infos):
                print(f"\n智能体 {i + 1}:")
                print(f"距离目标点的距离: {info['distance']:.2f}")
                print(f"是否发生碰撞: {info['collision']}")
                distance_sensor_data = info.get('distance_sensor_data', [])
                print(f"距离传感器数据: {distance_sensor_data}")

        test_reset_and_stop_functionality()

        # 关闭环境
        print("\n正在关闭环境...")
        # 由于 airsim_gym3.py 中未定义 close 方法，这里暂时忽略
        # env.close()
        print(f"\n测试完成！所有数据已保存到目录: {save_dir}")
        return True

    except Exception as e:
        print(f"\n测试失败---: {e}")
        return False


def save_environment_data(env, save_dir, step):
    """保存环境数据（彩色图、深度图和激光雷达）"""
    states = env.get_state()
    images_and_depths = env.get_image()  # 返回 (彩色图, 深度图) 元组列表
    lidar_data = env.get_lidar_data()

    # 保存图像和深度图数据
    for i, (rgb_img, depth_img) in enumerate(images_and_depths):
        # 保存彩色图
        rgb_path = os.path.join(save_dir, f"agent_{i + 1}_rgb_step_{step:02d}.png")
        cv2.imwrite(rgb_path, cv2.cvtColor(rgb_img, cv2.COLOR_RGB2BGR))

        # 保存深度图（16位PNG格式，单位为毫米）
        depth_path = os.path.join(save_dir, f"agent_{i + 1}_depth_step_{step:02d}.png")
        cv2.imwrite(depth_path, depth_img, [cv2.IMWRITE_PNG_COMPRESSION, 0])

        # 保存深度图的可视化版本（归一化到0-255）
        depth_vis = cv2.normalize(depth_img, None, 0, 255, cv2.NORM_MINMAX, dtype=cv2.CV_8U)
        depth_vis_path = os.path.join(save_dir, f"agent_{i + 1}_depth_vis_step_{step:02d}.png")
        cv2.imwrite(depth_vis_path, depth_vis)

    # 保存激光雷达数据
    for i, points in enumerate(lidar_data):
        lidar_path = os.path.join(save_dir, f"agent_{i + 1}_lidar_step_{step:02d}.txt")
        if points.size > 0:
            np.savetxt(lidar_path, points, fmt='%.6f', delimiter=' ')


def compare_initial_final_states(save_dir, total_steps):
    """比较初始和最终状态"""
    print("\n状态变化摘要:")

    for agent_id in range(1, 4):  # 三个智能体
        # 读取初始状态
        initial_lidar_file = os.path.join(save_dir, f"agent_{agent_id}_lidar_step_00.txt")
        initial_points = np.loadtxt(initial_lidar_file) if os.path.exists(initial_lidar_file) else np.array([])

        # 读取最终状态
        final_lidar_file = os.path.join(save_dir, f"agent_{agent_id}_lidar_step_{total_steps:02d}.txt")
        final_points = np.loadtxt(final_lidar_file) if os.path.exists(final_lidar_file) else np.array([])

        print(f"\n智能体 {agent_id}:")
        print(f"初始点云数量: {len(initial_points) if initial_points.size > 0 else 0}")
        print(f"最终点云数量: {len(final_points) if final_points.size > 0 else 0}")

        # 计算位置变化（简化版）
        if initial_points.size > 0 and final_points.size > 0:
            initial_center = np.mean(initial_points[:, :2], axis=0)  # XY平面中心
            final_center = np.mean(final_points[:, :2], axis=0)
            displacement = np.linalg.norm(final_center - initial_center)
            print(f"估计位移: {displacement:.2f} 米")


def test_reset_and_stop_functionality():
    """测试AirSimMultiAgentEnv的重置和停止功能"""
    try:
        print("===== 测试环境重置和停止功能 =====")
        # 创建环境实例
        print("正在创建环境实例...")
        env = AirSimMultiAgentEnv()

        # 测试重置功能
        print("\n===== 测试环境重置 =====")
        for i in range(3):  # 测试3次重置
            print(f"\n执行第 {i + 1} 次重置测试")
            target_pos = [10 + i * 5, 5 + i * 2, 0]  # 每次使用不同的目标位置
            print(f"设置目标位置: {target_pos}")

            # 重置环境
            env.reset(target_pos)
            print("环境重置成功")

            # 获取初始状态信息
            infos = env.get_info(target_pos)
            if infos:
                for j, info in enumerate(infos):
                    print(f"智能体 {j + 1} 初始距离目标点: {info['distance']:.2f}")
                    distance_sensor_data = info.get('distance_sensor_data', [])
                    print(f"智能体 {j + 1} 初始距离传感器数据: {distance_sensor_data}")

            # 执行几个步骤的随机动作
            for step in range(5):
                random_actions = [
                    [np.random.uniform(0.5, 0.9),  # 速度
                     np.random.uniform(-0.8, 0.8),  # 转向
                     0.0]  # 刹车
                    for _ in range(3)  # 为每个智能体生成随机动作
                ]
                states, rewards, dones, infos = env.step(random_actions)
                print(f"步骤 {step + 1}: 奖励 = {rewards}, 完成 = {dones}")
                for j, info in enumerate(infos['infos']):
                    distance_sensor_data = info.get('distance_sensor_data', [])
                    print(f"智能体 {j + 1} 步骤 {step + 1} 距离传感器数据: {distance_sensor_data}")

            # 检查重置后的状态是否正确
            infos_after_reset = env.get_info(target_pos)
            if infos_after_reset:
                print("\n重置后状态检查:")
                for j, info in enumerate(infos_after_reset):
                    print(f"智能体 {j + 1} 距离目标点: {info['distance']:.2f}, 碰撞状态: {info['collision']}")
                    distance_sensor_data = info.get('distance_sensor_data', [])
                    print(f"智能体 {j + 1} 重置后距离传感器数据: {distance_sensor_data}")

        # 测试停止功能
        print("\n===== 测试环境停止 =====")
        print("正在停止环境...")
        # 假设AirSimMultiAgentEnv有一个close方法用于停止环境
        if hasattr(env, 'close') and callable(getattr(env, 'close')):
            env.close()
            print("环境已成功停止")
        else:
            print("警告: AirSimMultiAgentEnv类中未定义close方法")

        # 尝试在停止后调用方法，验证是否抛出适当的错误
        try:
            print("\n尝试在环境停止后获取信息...")
            env.get_info([0, 0, 0])
        except Exception as e:
            print(f"预期错误: {e}")
        else:
            print("警告: 在环境停止后仍能成功调用方法")

        print("\n===== 测试完成 =====")
        return True

    except Exception as e:
        print(f"\n测试失败: {e}")
        return False


# 运行测试
if __name__ == "__main__":
    test_airsim_multi_agent_env()